package rabbitmq.publisher.confirms;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import rabbitmq.constant.Constants;

import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.SortedSet;
import java.util.TreeSet;


public class PublisherConfirms {

    private static final Integer MESSAGE_COUNT = 200;

    static Connection createConnection() throws Exception {
        ConnectionFactory connectionFactory = new ConnectionFactory();
        connectionFactory.setHost(Constants.HOST);
        connectionFactory.setPort(Constants.PORT);
        connectionFactory.setUsername(Constants.USER_NAME);
        connectionFactory.setPassword(Constants.PASSWORD);
        connectionFactory.setVirtualHost(Constants.VIRTUAL_HOST);

        return connectionFactory.newConnection();
    }

    public static void main(String[] args) throws Exception {

        //单独确认
        Individually();

        //批量确认
        Batches();

        //异步确认
        Asynchronously();
    }


    private static void Individually() throws Exception {
        try(Connection connection = createConnection()) {
            //1.开启信道
            Channel channel = connection.createChannel();
            //2.设置信道为 confirm 模式
            channel.confirmSelect();
            //3.声明队列(交换机为默认交换机)
            channel.queueDeclare(Constants.PUBLISH_CONFIRMS_QUEUE1, true, false, false, null);
            //4.发送消息,并等待确认
            long start = System.currentTimeMillis();
            for(int i = 0; i <MESSAGE_COUNT; i++){
                String msg = "publisher confirm " + i;
                channel.basicPublish("", Constants.PUBLISH_CONFIRMS_QUEUE1, null, msg.getBytes());
                //等待确认
                channel.waitForConfirmsOrDie(4000);
            }
            long end = System.currentTimeMillis();
            System.out.printf("单独确认 - 消息条数 : %d, 耗时 : %d ms \n", MESSAGE_COUNT, end - start);
        }
    }

    private static void Batches() throws Exception {
        try(Connection connection = createConnection()){
            //1.开启信道
            Channel channel = connection.createChannel();
            //2.设置信道为 confirm 模式
            channel.confirmSelect();
            //3.声明队列(交换机为默认交换机)
            channel.queueDeclare(Constants.PUBLISH_CONFIRMS_QUEUE2, true, false, false, null);
            //4.发送消息,并等待确认
            int batchSize = 100;
            int outstandingMessageCount = 0;
            long start = System.currentTimeMillis();
            for(int i = 0; i < MESSAGE_COUNT; i++){
                String msg = "publisher confirm " + i;
                channel.basicPublish("", Constants.PUBLISH_CONFIRMS_QUEUE2, null, msg.getBytes());
                outstandingMessageCount++;
                if(outstandingMessageCount == batchSize){
                    channel.waitForConfirmsOrDie(4000);
                    outstandingMessageCount = 0;
                }
            }
            if(outstandingMessageCount > 0){
                channel.waitForConfirmsOrDie(4000);
            }
            long end = System.currentTimeMillis();
            System.out.printf("批量确认 - 消息条数 : %d, 耗时 : %d ms \n", MESSAGE_COUNT, end - start);
        }
    }

    private static void Asynchronously() throws Exception {
        try(Connection connection = createConnection()) {
            //1.开启信道
            Channel channel = connection.createChannel();
            //2.设置信道为 confirm 模式
            channel.confirmSelect();
            //3.声明队列(交换机为默认交换机)
            channel.queueDeclare(Constants.PUBLISH_CONFIRMS_QUEUE3, true, false, false, null);
            //4.监听 confirm
            //集合中存储的是未确认的消息ID
            long start = System.currentTimeMillis();
            SortedSet<Long> confirmSeqNo = Collections.synchronizedSortedSet(new TreeSet<>());

            channel.addConfirmListener(new ConfirmListener() {
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    if(multiple){
                        confirmSeqNo.headSet(deliveryTag + 1).clear();
                    }else{
                        confirmSeqNo.remove(deliveryTag);
                    }
                }

                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                    if(multiple){
                        confirmSeqNo.headSet(deliveryTag + 1).clear();
                    }else{
                        confirmSeqNo.remove(deliveryTag);
                    }
                    //业务需要根据实际场景进行处理,比如重发
                }
            });
            //5.发送消息
            for(int i = 0; i <MESSAGE_COUNT; i++){
                String msg = "publisher confirm " + i;
                long seqNo = channel.getNextPublishSeqNo();
                channel.basicPublish("", Constants.PUBLISH_CONFIRMS_QUEUE3, null, msg.getBytes());
                confirmSeqNo.add(seqNo);
            }
            while(!confirmSeqNo.isEmpty()){
                Thread.sleep(10);
            }
            long end = System.currentTimeMillis();
            System.out.printf("异步确认 - 消息条数 : %d, 耗时 : %d ms \n", MESSAGE_COUNT, end - start);
        }
    }
}
